from torch._C import device
from NN.bnn_tree import BNNTree
from util.visualize_nn import *
from NN.ResNet import ResNet18_WordNet_AUG

if __name__ == "__main__":
    # 可视化 resnet-wordnet-augmentation
    # visualize_RNWN_AUG()

    # 测试 mistgpu 训练结果
    # from torchvision.transforms import ToTensor
    # from torchvision import datasets, transforms
    # from torch.utils.data import DataLoader

    # transform_test = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    #     ])
    # test_data = datasets.CIFAR10('./data', train=False,
    #                    transform=transform_test)
    # batch_size = 128
    # test_dataloader = DataLoader(test_data, batch_size=batch_size) # 10000

    # model_rnwn = rnwn_rule_aug_ = ResNet18_WordNet_AUG("resnet18-cifar10-rule-aug-nopretrained", 3, 10, "cpu")
    # model_rnwn.pre_train(85)
    # model_rnwn.nn_test(test_dataloader, model_rnwn.loss_fn) # v1: epoch-85-88.3%; v2: epoch-86 87.5%
    # model_rnwn.draw_process()
    # model_rnwn.save_process("resnet18-cifar10-rule-aug-nopretrained-v2", os.path.join("output", f"{model_rnwn.model_note}_v2.png"))

    
    # 分析 rule(9)


    # 二分神经网络树 ==================================================================
    from util.bnnt_dataset import *
    from NN.bnn_tree import *
    from torch.optim import *
    from torchvision import datasets
    from NN.bnnt_node import *

    # mnist
    # mnistdata = BNNTDATA(batch_size=128)
    # bnntree = BNNTree(10, 1, mnistdata, 5)
    # # print(bnntree.test_acc())
    # bnntree.set_optim_lr(Adam, 1e-1)
    # bnntree.train(bnntree.tree, 0, 9)
    # print(bnntree.test_acc()) # 97.51%
    
    # cifar10
    from combination_train.bnnt_train import bnnt_train_cifar10
    bnnt_train_cifar10()

    pass